package mr

import (
	"log"
	"net"
	"net/http"
	"net/rpc"
	"os"
	"strconv"
	"time"
)

type Coordinator struct {
	// Your definitions here.
	nMap    int
	nReduce int
	phase   TaskPhase
	files   []string
	tasks   []Task

	taskrequestCh  chan taskRequestMsg
	taskresponseCh chan taskResponseMsg
	doneCh         chan struct{}
}

type taskRequestMsg struct {
	reply *TaskReqReply
	ok    chan struct{}
}

type taskResponseMsg struct {
	response *ResponseArg
	ok       chan struct{}
}

// Your code here -- RPC handlers for the worker to call.

// an example RPC handler.
//
// the RPC argument and reply types are defined in rpc.go.
func (c *Coordinator) Example(args *ExampleArgs, reply *ExampleReply) error {
	reply.Y = args.X + 1
	return nil
}

func (c *Coordinator) TaskRequest(args *TaskReqArg, reply *TaskReqReply) error {
	msg := taskRequestMsg{reply: reply, ok: make(chan struct{})}
	c.taskrequestCh <- msg
	<-msg.ok
	return nil
}

func (c *Coordinator) TaskReport(args *ResponseArg, reply *ResponseReply) error {
	msg := taskResponseMsg{response: args, ok: make(chan struct{})}
	c.taskresponseCh <- msg
	<-msg.ok
	return nil
}

// start a thread that listens for RPCs from worker.go
func (c *Coordinator) server() {
	rpc.Register(c)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := coordinatorSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
func (c *Coordinator) Done() bool {
	// Your code here.
	select {
	case <-c.doneCh:
		return true
	default:
		return false
	}

}

func (c *Coordinator) initMapPhase() {
	c.phase = MapPhase
	c.tasks = make([]Task, c.nMap)
	for i := range c.tasks {
		c.tasks[i] = Task{
			Status:    Waiting,
			StartTime: 0,
			FileId:    i,
			FileName:  c.files[i],
		}
	}
}

func (c *Coordinator) initReducePhase() {
	c.phase = ReducePhase
	c.tasks = make([]Task, c.nReduce)
	for i := range c.tasks {
		c.tasks[i] = Task{
			Status:    Waiting,
			StartTime: 0,
			FileId:    i,
			FileName:  "",
		}
	}
}

func (c *Coordinator) checkTaskStatus() bool {
	done := false
	for i := range c.tasks {
		if c.tasks[i].Status != Done {
			return false
		}
	}
	if c.phase == ReducePhase {
		done = true
		c.doneCh <- struct{}{}
	} else {
		c.initReducePhase()
	}
	return done
}

func (c *Coordinator) getWaitingTask() (res int) {
	res = -1
	for i := range c.tasks {
		if c.tasks[i].Status == Done {
			continue
		} else if c.tasks[i].Status == Waiting {
			res = i
			return
		} else if time.Now().Unix()-c.tasks[i].StartTime >= 10 { // Task overtime, redo
			res = i
			return
		} else {
			res = -2
		}
	}
	return
}

func (c *Coordinator) cleanIntermediateFile(id int) {
	for i := 0; i < c.nMap; i++ {
		fname := "mr-" + strconv.Itoa(i) + "-" + strconv.Itoa(id)
		if err := os.RemoveAll(fname); err != nil {
			log.Printf("delete intermediate file %v failed, err: %v\n", fname, err)
		}
	}
}

func (c *Coordinator) schedule() {
	c.initMapPhase()
	for {
		select {
		case msg := <-c.taskrequestCh:
			undo := c.getWaitingTask()
			if undo >= 0 {
				msg.reply.Id = undo
				if c.phase == MapPhase {
					msg.reply.Type = MapTask
					msg.reply.FileName = c.files[undo]
				} else {
					msg.reply.Type = ReduceTask
				}
				msg.reply.NReduce = c.nReduce
				msg.reply.NMap = c.nMap
				c.tasks[undo].Status = Doing
				c.tasks[undo].StartTime = time.Now().Unix()
			} else if undo == -1 && c.phase == ReducePhase { // AllTaskDone
				msg.reply.Type = ExistTask
			} else { // cannot find task can do now, let worker wait
				msg.reply.Type = WaitTask
			}
			msg.ok <- struct{}{}
		case msg := <-c.taskresponseCh:
			// task may be timeout
			if c.tasks[msg.response.Id].Status == Doing {
				c.tasks[msg.response.Id].Status = Done
				// clean the intermediate output
				if c.phase == ReducePhase {
					c.cleanIntermediateFile(msg.response.Id)
				}
			} else {
				continue
			}
			// check if map is over or all task is over
			if c.checkTaskStatus() {
				return
			}
			msg.ok <- struct{}{}
		}
	}
}

// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
func MakeCoordinator(files []string, nReduce int) *Coordinator {
	c := Coordinator{}
	c.nMap = len(files)
	c.files = files
	c.nReduce = nReduce
	c.doneCh = make(chan struct{}, 1)
	c.taskrequestCh = make(chan taskRequestMsg)
	c.taskresponseCh = make(chan taskResponseMsg)
	// Your code here.
	go c.schedule()
	c.server()
	return &c
}
